
"""
数据集过滤
"""
from pydantic import Field
from core.llm.base_analyze import BaseAnalyzer
from core import prompt_dir
from typing import List
from models.base import BaseData
from models.aiquery import AiQueryData, TableData
from loguru import logger


class DatasetFilterAnalyzer(BaseAnalyzer):

    def __init__(self):
        super().__init__(prompt_file=prompt_dir / '数据筛选.md')

    def invoke(
            self,
            datasets: List[BaseData] = Field(description="数据表"),
            query: str = Field(description="用户问题"),
    ):
        descriptions = dict()
        for idx, dataset in enumerate(datasets):
            if dataset.type == 'text':
                descriptions[idx] = dataset.model_dump()
            elif isinstance(dataset, AiQueryData) and isinstance(dataset.data, TableData):
                data = dataset.data.model_dump()
                data.pop('columns', None)
                data.pop('rows', None)
                descriptions[idx] = data

        prompt = self.template_prompt.format(
            inputs={
                "descriptions": descriptions,
                "query": query
            },
            remove_template_variables=True
        )
        indexes = self.analyze(prompt)['indexes']
        if not isinstance(indexes, list):
            return []

        filtered_datasets = []
        for index in indexes:
            if isinstance(index, str) and index.isdigit():
                index = int(index)
            if not isinstance(index, int):
                continue
            if 0 <= index < len(datasets):
                filtered_datasets.append(datasets[index])
        return filtered_datasets
